Fix mesh axis bugs & Fix Wan-VACE transformer and pipeline to align with other Wan models and pipelines#390
Merged
copybara-service[bot] merged 1 commit intomainfrom Apr 29, 2026
Conversation
ninatu
changed the title
ninatu
changed the title
…ith other Wan models and pipeline * Remove an unnecessary `vae_spatial_axis_name` assignment to `vae_mesh` in `wan_pipeline.py`, that fixes `AttributeError: cannot assign to field 'vae_spatial_axis_name'`. * Adopts the shared `_create_common_components` method (used by `WanPipeline2_1`, `WanPipeline2_2`, etc.) into `_load_and_init` of `VaceWanPipeline2_1`. This align initialization and resolves mesh axis handling issues during `VaceWanPipeline2_1` initialization introduced by PR #359. * Adopts new parameters `mask_padding_tokens`, `enable_jax_named_scopes`, `use_base2_exp`, `use_experimental_scheduler` into `transformer_wan_vace.py` and `wan_vace_pipeline_2_1.py` (following `transformer_wan.py` and `wan_pipeline.py` respectively) that resolves attention errors (introduced by PR #359) and align Wan-VACE transformer to Wan transformer. * Removes the `load_common_components` argument from `from_pretrained` and `from_checkpoint` in `VaeWanPipeline2_1` for interface consistency.
ninatu
force-pushed
the
ninatu/wan-vace-mesh-init-bugs
branch
from
10ade80 to
fa877ab
Compare
entrpn
reviewed
Apr 28, 2026
| vae_devices_array = flat_devices.reshape(total_devices // vae_spatial, vae_spatial) | ||
|
|
||
| vae_mesh = Mesh(vae_devices_array, ("redundant", "vae_spatial")) | ||
| vae_mesh.vae_spatial_axis_name = "vae_spatial" |
Collaborator
There was a problem hiding this comment.
@eltsai can you check the significance of this change.
Collaborator
Author
There was a problem hiding this comment.
hey @eltsai! just a kind reminder to take a look :)
Collaborator
There was a problem hiding this comment.
yes it's safe to remove, in WAN VAE sharding we didn't read from vae_mesh
eltsai
added a commit
that referenced
this pull request
Apr 29, 2026
entrpn
approved these changes
Apr 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Remove an unnecessary
vae_spatial_axis_nameassignment tovae_meshinwan_pipeline.py, that fixesAttributeError: cannot assign to field 'vae_spatial_axis_name'.Adopts the shared
_create_common_componentsmethod (used byWanPipeline2_1,WanPipeline2_2, etc.) into_load_and_initofVaceWanPipeline2_1. This align initialization and resolves mesh axis handling issues duringVaceWanPipeline2_1initialization introduced by PR Ring attention integration and other optimizations #359.Adopts new parameters
mask_padding_tokens,enable_jax_named_scopes,use_base2_exp,use_experimental_schedulerintotransformer_wan_vace.pyandwan_vace_pipeline_2_1.py(followingtransformer_wan.pyandwan_pipeline.pyrespectively) that align Wan-VACE transformer to Wan transformer.Removes the
load_common_componentsargument fromfrom_pretrainedandfrom_checkpointinVaceWanPipeline2_1for interface consistency.